package markov;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;

public class WordBigram {
    ArrayList<String[]> arr;
    HashMap<String, HashMap<String,Integer>> posCount;
    public HashMap<String, HashMap<String,Double>> wordProbability;
    public HashMap<String, HashMap<String,Double>> wordProbabilityWordIndex;
    HashMap<String, Integer> posTotal;
    public Boolean PRING_DEBUG_INFO;
    public WordBigram(ArrayList<String[]> arr){
	this.arr = arr;
	posCount = new HashMap<String, HashMap<String,Integer>>();
	wordProbability = new HashMap<String, HashMap<String,Double>>();
	posTotal = new HashMap<String, Integer>();
	PRING_DEBUG_INFO = false;
	wordProbabilityWordIndex = new HashMap<String, HashMap<String,Double>>();
    }

    public void run(){
	for( String[] cur : arr){
	    if(posTotal.containsKey(cur[1]))
		posTotal.put(cur[1], posTotal.get(cur[1])+1); //increment total count
	    else
		posTotal.put(cur[1], 1);

	    if(posCount.containsKey(cur[1])){ //word
		//TODO contains type key
		int val = 0;

		if(posCount.get(cur[1]).containsKey(cur[0]))
		    val = posCount.get(cur[1]).get(cur[0]);
		else
		    posCount.get(cur[1]).put(cur[0], 1);
		posCount.get(cur[1]).put(cur[0], ++val);
	    }else{
		final String tmp = cur[0];
		posCount.put(cur[1], new HashMap<String, Integer>(){{put(tmp ,1);}});
	    }
	}
	generateProbability();

	if(PRING_DEBUG_INFO) printCount();
    }

    public void generateProbability(){
	for(final Map.Entry<String, HashMap<String,Integer>> a : posCount.entrySet()){
	    for(final Map.Entry<String,Integer> b : a.getValue().entrySet()){
		final double prob = b.getValue() / (double) posTotal.get(a.getKey());

		if(wordProbability.containsKey(a.getKey())){
		    wordProbability.get(a.getKey()).put(b.getKey(), prob);
		}
		else{
		    wordProbability.put(a.getKey(), new HashMap(){{put(b.getKey(),prob);}});	
		}
		
		if(wordProbabilityWordIndex.containsKey(b.getKey()))
		    wordProbabilityWordIndex.get(b.getKey()).put(a.getKey(), prob);
		else 
		    wordProbabilityWordIndex.put(b.getKey(), new HashMap(){{put(a.getKey(),prob);}});
	    }
	}
    }
    public void printCount(){
	for(Map.Entry<String, HashMap<String,Double>> pos1 : wordProbability.entrySet()){
	    for(Entry<String, Double> pos2 : pos1.getValue().entrySet()){
		System.out.println(pos1.getKey() + "\t" + pos2.getKey() + "\t" + (pos2.getValue()));
	    }

	}
    }

}
